src/utils/arrow/mod.rs (1)

144-183: Well-implemented function for adding custom fields to record batches

This function is well-structured and handles the addition of custom fields properly:

• Correctly preserves the original record batch data
• Inserts timestamp at the beginning
• Sorts custom fields by key for consistent ordering
• Handles field type definitions appropriately

One minor optimization suggestion would be to pre-allocate the vectors with the correct capacity.

You could optimize vector allocations by pre-calculating the final size:

-    let mut fields: Vec<Field> = schema.fields().iter().map(|f| f.as_ref().clone()).collect();
-    let mut columns: Vec<ArrayRef> = rb.columns().to_vec();
+    let final_size = schema.fields().len() + 1 + p_custom_fields.len();
+    let mut fields: Vec<Field> = Vec::with_capacity(final_size);
+    let mut columns: Vec<ArrayRef> = Vec::with_capacity(final_size);
+    
+    // Add existing fields and columns
+    fields.extend(schema.fields().iter().map(|f| f.as_ref().clone()));
+    columns.extend(rb.columns().to_vec());

src/handlers/http/modal/utils/ingest_utils.rs (1)

163-168: Consider adding validation for custom header keys.

While the implementation correctly extracts headers with the x-p- prefix, consider adding validation to ensure the trimmed keys follow any specific naming restrictions your schema might have.

 if header_name.starts_with("x-p-") && !IGNORE_HEADERS.contains(&header_name) {
     if let Ok(value) = header_value.to_str() {
         let key = header_name.trim_start_matches("x-p-");
+        // Validate key format - alphanumeric and underscore only
+        if key.chars().all(|c| c.is_alphanumeric() || c == '_') {
             p_custom_fields.insert(key.to_string(), value.to_string());
+        }
     }
 }

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/utils/arrow/mod.rs (1)

204-270: 💡 Verification agent

❓ Verification inconclusive

Missing tests for the new add_parseable_fields function.

While there are tests for other functions in this module, there are no tests specifically for the newly added add_parseable_fields function. Consider adding tests that verify:

1. Custom fields are correctly added to the RecordBatch
2. The timestamp is correctly inserted
3. Error cases are properly handled
4. Behavior when the input RecordBatch is empty
5. Behavior with empty custom fields

---

🏁 Script executed:

#!/bin/bash
# Check if there are any tests for the add_parseable_fields function in the codebase
rg -l "test.*add_parseable_fields" --type rust

Length of output: 48

---

Actionable: Add tests for the new add_parseable_fields function

After verifying the repository, there are no tests referencing add_parseable_fields. Please add tests covering the following aspects:

• Custom fields are correctly added to the RecordBatch.
• The timestamp is correctly inserted.
• Error cases are properly handled.
• Behavior when the input RecordBatch is empty.
• Behavior with empty custom fields.

src/utils/arrow/mod.rs (1)

144-192: The implementation correctly adds custom fields to RecordBatches.

The add_parseable_fields function properly handles the addition of a timestamp and custom fields to the RecordBatch. The sorting of keys ensures consistent field order, and the implementation correctly maintains the Arrow memory model with appropriate use of Arc.

Consider these improvements:

1. Add tests for this new function to verify behavior with various inputs
2. Check if a timestamp field already exists in the schema before inserting it
3. Add documentation for possible errors that could be returned

 pub fn add_parseable_fields(
     rb: RecordBatch,
     p_timestamp: DateTime<Utc>,
     p_custom_fields: &HashMap<String, String>,
 ) -> Result<RecordBatch, ArrowError> {
     // Return Result for proper error handling
+    // Possible errors:
+    // - Schema mismatch if fields cannot be combined
+    // - Memory allocation errors for large batches

     // Add custom fields in sorted order
     let mut sorted_keys: Vec<&String> = p_custom_fields.keys().collect();
     sorted_keys.sort();

     let schema = rb.schema();
     let row_count = rb.num_rows();

+    // Check if timestamp field already exists
+    let timestamp_field_exists = schema.fields().iter().any(|f| f.name() == DEFAULT_TIMESTAMP_KEY);
+
     let mut fields = schema
         .fields()
         .iter()
         .map(|f| f.as_ref().clone())
         .collect_vec();
-    fields.insert(
-        0,
-        Field::new(
-            DEFAULT_TIMESTAMP_KEY,
-            DataType::Timestamp(TimeUnit::Millisecond, None),
-            true,
-        ),
-    );
+    if !timestamp_field_exists {
+        fields.insert(
+            0,
+            Field::new(
+                DEFAULT_TIMESTAMP_KEY,
+                DataType::Timestamp(TimeUnit::Millisecond, None),
+                true,
+            ),
+        );
+    }
     fields.extend(
         sorted_keys
             .iter()
             .map(|k| Field::new(*k, DataType::Utf8, true)),
     );

     let mut columns = rb.columns().iter().map(Arc::clone).collect_vec();
-    columns.insert(
-        0,
-        Arc::new(get_timestamp_array(p_timestamp, row_count)) as ArrayRef,
-    );
+    if !timestamp_field_exists {
+        columns.insert(
+            0,
+            Arc::new(get_timestamp_array(p_timestamp, row_count)) as ArrayRef,
+        );
+    }
     columns.extend(sorted_keys.iter().map(|k| {
         let value = p_custom_fields.get(*k).unwrap();
         Arc::new(StringArray::from_iter_values(
             std::iter::repeat(value).take(row_count),
         )) as ArrayRef
     }));

     // Create the new schema and batch
     let new_schema = Arc::new(Schema::new(fields));
     RecordBatch::try_new(new_schema, columns)
 }

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/handlers/http/modal/utils/ingest_utils.rs (3)

146-156: Good implementation for extracting HTTP headers.

The function correctly extracts the User-Agent and source IP from the request. Consider adding error logging for cases where these values cannot be extracted.

 pub fn get_custom_fields_from_header(req: HttpRequest) -> HashMap<String, String> {
     let user_agent = req
         .headers()
         .get("User-Agent")
         .and_then(|a| a.to_str().ok())
-        .unwrap_or_default();
+        .unwrap_or_else(|| {
+            log::debug!("User-Agent header not found or not valid UTF-8");
+            ""
+        });

     let conn = req.connection_info().clone();

-    let source_ip = conn.realip_remote_addr().unwrap_or_default();
+    let source_ip = conn.realip_remote_addr().unwrap_or_else(|| {
+        log::debug!("Source IP not found in connection info");
+        ""
+    });

---

157-158: Consider conditional inclusion of empty values.

The current implementation always adds User-Agent and source IP fields, even when they're empty. You might want to add them only when they contain non-empty values.

-    p_custom_fields.insert(USER_AGENT_KEY.to_string(), user_agent.to_string());
-    p_custom_fields.insert(SOURCE_IP_KEY.to_string(), source_ip.to_string());
+    if !user_agent.is_empty() {
+        p_custom_fields.insert(USER_AGENT_KEY.to_string(), user_agent.to_string());
+    }
+    if !source_ip.is_empty() {
+        p_custom_fields.insert(SOURCE_IP_KEY.to_string(), source_ip.to_string());
+    }

---

161-169: Check for duplicate header names.

If multiple headers with the same name (after stripping "x-p-") are present, the last one will overwrite previous values. Consider logging a warning when this happens, or implement a strategy to handle duplicates (e.g., comma-separated values).

             if let Ok(value) = header_value.to_str() {
                 let key = header_name.trim_start_matches("x-p-");
-                p_custom_fields.insert(key.to_string(), value.to_string());
+                if p_custom_fields.contains_key(key) {
+                    log::warn!("Duplicate custom field key '{}', overwriting previous value", key);
+                }
+                p_custom_fields.insert(key.to_string(), value.to_string());
             }

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/utils/arrow/mod.rs (2)

146-194: Consider adding tests for this new function

Given the importance of this function for the custom fields feature, it would be beneficial to add unit tests that verify:

• Custom fields are properly added
• Fields appear in the expected order
• The row count is preserved
• Edge cases (empty HashMap, large number of fields, etc.)

---

185-186: Potential unwrap() risk

The unwrap() call on line 185 could potentially panic if the key doesn't exist in the HashMap. While this appears safe since the keys are directly derived from the HashMap itself in the loop, consider using .expect() with a descriptive message for better error reporting in case of future modifications.

-        let value = p_custom_fields.get(*k).unwrap();
+        let value = p_custom_fields.get(*k).expect("Key must exist in HashMap");

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/handlers/http/modal/utils/ingest_utils.rs (2)

147-169: 🛠️ Refactor suggestion

Add validation for empty keys after stripping prefix.

The function doesn't validate if a key is empty after trimming the "x-p-" prefix. If a header with name "x-p-" (with nothing after) is provided, it would result in an empty key in the HashMap.

        if header_name.starts_with("x-p-") && !IGNORE_HEADERS.contains(&header_name) {
            if let Ok(value) = header_value.to_str() {
                let key = header_name.trim_start_matches("x-p-");
-               p_custom_fields.insert(key.to_string(), value.to_string());
+               if !key.is_empty() {
+                   p_custom_fields.insert(key.to_string(), value.to_string());
+               } else {
+                   log::warn!("Ignoring header with empty key after prefix: {}", header_name);
+               }
            }
        }

---

147-179: 🛠️ Refactor suggestion

Implement sanitization for custom field values.

Header values are inserted directly into the custom fields HashMap without any sanitization. This could potentially lead to injection attacks or issues with malformed data affecting downstream processing.

Consider adding a sanitization function for the header values before insertion:

// Add this utility function
+ fn sanitize_header_value(value: &str) -> String {
+     // Implement appropriate sanitization logic
+     // This could include escaping special characters,
+     // removing control characters, limiting length, etc.
+     value.chars()
+         .filter(|&c| !c.is_control())
+         .take(1000)  // Limit size
+         .collect()
+ }

// Then use it when inserting values
-    p_custom_fields.insert(USER_AGENT_KEY.to_string(), user_agent.to_string());
+    p_custom_fields.insert(USER_AGENT_KEY.to_string(), sanitize_header_value(user_agent));

-    p_custom_fields.insert(SOURCE_IP_KEY.to_string(), source_ip.to_string());
+    p_custom_fields.insert(SOURCE_IP_KEY.to_string(), sanitize_header_value(source_ip));

// And in the header processing loop
-                p_custom_fields.insert(key.to_string(), value.to_string());
+                p_custom_fields.insert(key.to_string(), sanitize_header_value(value));

// And for the format key
-                p_custom_fields.insert(FORMAT_KEY.to_string(), value.to_string());
+                p_custom_fields.insert(FORMAT_KEY.to_string(), sanitize_header_value(value));

src/handlers/http/modal/utils/ingest_utils.rs (2)

164-168: Consider case-insensitive header name matching.

HTTP headers are case-insensitive by specification, but the check header_name.starts_with("x-p-") is case-sensitive. This could cause issues if clients send headers with different casing (e.g., "X-P-Environment").

-        if header_name.starts_with("x-p-") && !IGNORE_HEADERS.contains(&header_name) {
+        let header_name_lower = header_name.to_lowercase();
+        if header_name_lower.starts_with("x-p-") && !IGNORE_HEADERS.iter().any(|&h| h.to_lowercase() == header_name_lower) {
             if let Ok(value) = header_value.to_str() {
-                let key = header_name.trim_start_matches("x-p-");
+                let key = header_name_lower.trim_start_matches("x-p-");
                 p_custom_fields.insert(key.to_string(), value.to_string());
             }
         }

---

181-224: Tests are well-structured but could be more comprehensive.

The tests cover basic functionality but could be expanded to include edge cases:

• Test with malformed header values
• Test with empty keys after stripping "x-p-"
• Test with case variations (e.g., X-P-Environment vs x-p-environment)
• Test with duplicate keys (standard + custom headers mapping to same key)

Consider adding the following test case:

#[test]
fn test_get_custom_fields_from_header_with_case_insensitive_headers() {
    let req = TestRequest::default()
        .insert_header((USER_AGENT, "TestUserAgent"))
        .insert_header(("X-P-Environment", "dev"))  // Uppercase X-P
        .to_http_request();

    let custom_fields = get_custom_fields_from_header(req);

    assert_eq!(custom_fields.get(USER_AGENT_KEY).unwrap(), "TestUserAgent");
    assert_eq!(custom_fields.get("environment").unwrap(), "dev");
}

#[test]
fn test_get_custom_fields_from_header_with_empty_key() {
    let req = TestRequest::default()
        .insert_header((USER_AGENT, "TestUserAgent"))
        .insert_header(("x-p-", "value"))  // Empty key after prefix
        .to_http_request();

    let custom_fields = get_custom_fields_from_header(req);

    assert_eq!(custom_fields.get(USER_AGENT_KEY).unwrap(), "TestUserAgent");
    assert!(!custom_fields.contains_key(""));  // Empty key should be ignored
}

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/utils/arrow/mod.rs (1)

184-195: Consider adding type detection for custom field values

Currently, all custom fields are added as UTF8 string types. For improved data analytics, you might want to consider detecting and converting values to appropriate Arrow types (numbers, booleans, etc.) when possible.

 for key in sorted_keys {
     if !field_names.contains(key) {
-        fields.push(Field::new(key, DataType::Utf8, true));
+        // Try to parse the value into more specific types
+        let value = p_custom_fields.get(key).unwrap();
+        let data_type = if let Ok(_) = value.parse::<i64>() {
+            DataType::Int64
+        } else if let Ok(_) = value.parse::<f64>() {
+            DataType::Float64
+        } else if value == "true" || value == "false" {
+            DataType::Boolean
+        } else {
+            DataType::Utf8
+        };
+        fields.push(Field::new(key, data_type, true));
         field_names.insert(key.to_string());

         let value = p_custom_fields.get(key).unwrap();
-        columns.push(Arc::new(StringArray::from_iter_values(
-            std::iter::repeat(value).take(row_count),
-        )) as ArrayRef);
+        // Create appropriate array based on detected type
+        let array: ArrayRef = if let DataType::Int64 = data_type {
+            Arc::new(arrow_array::Int64Array::from_iter_values(
+                std::iter::repeat(value.parse::<i64>().unwrap()).take(row_count)
+            ))
+        } else if let DataType::Float64 = data_type {
+            Arc::new(arrow_array::Float64Array::from_iter_values(
+                std::iter::repeat(value.parse::<f64>().unwrap()).take(row_count)
+            ))
+        } else if let DataType::Boolean = data_type {
+            Arc::new(arrow_array::BooleanArray::from_iter_values(
+                std::iter::repeat(value == "true").take(row_count)
+            ))
+        } else {
+            Arc::new(StringArray::from_iter_values(
+                std::iter::repeat(value).take(row_count)
+            ))
+        };
+        columns.push(array);
     }
 }

src/handlers/http/modal/utils/ingest_utils.rs (1)

164-175: Implement field name validation for Arrow schema compatibility

While you've correctly handled empty keys after trimming the "x-p-" prefix, there's no validation that the resulting key names are valid for Arrow schema fields. Consider adding validation to ensure key names don't contain characters that might be problematic in Arrow field names.

 if header_name.starts_with("x-p-") && !IGNORE_HEADERS.contains(&header_name) {
     if let Ok(value) = header_value.to_str() {
         let key = header_name.trim_start_matches("x-p-");
         if !key.is_empty() {
-            p_custom_fields.insert(key.to_string(), value.to_string());
+            // Validate that the key name is valid for Arrow schema
+            if is_valid_arrow_field_name(key) {
+                p_custom_fields.insert(key.to_string(), value.to_string());
+            } else {
+                warn!(
+                    "Ignoring header with invalid Arrow field name: {}",
+                    header_name
+                );
+            }
         } else {
             warn!(
                 "Ignoring header with empty key after prefix: {}",
                 header_name
             );
         }
     }
 }

 // Add a helper function to validate field names
+fn is_valid_arrow_field_name(name: &str) -> bool {
+    // Implement validation logic based on Arrow's requirements
+    // For example, disallow spaces and special characters
+    !name.contains(|c: char| c.is_whitespace() || ".,;:!@#$%^&*()+=[]{}|\\<>/".contains(c))
+}

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

src/handlers/http/modal/utils/ingest_utils.rs (2)

149-207: Well-implemented extraction of custom fields with proper validation.

The implementation correctly:

• Extracts user agent and source IP information
• Processes headers with "x-p-" prefix
• Validates non-empty keys
• Implements length limits with truncation
• Handles special cases (LOG_SOURCE_KEY)

Consider adding a brief documentation comment explaining the function's purpose, expected header format, and how it handles the special case for LOG_SOURCE_KEY.

+/**
+ * Extracts custom fields from HTTP request headers.
+ * 
+ * This function:
+ * - Always includes User-Agent and source IP as default fields
+ * - Processes headers with "x-p-" prefix as custom fields
+ * - Handles LOG_SOURCE_KEY specially by adding it as a FORMAT_KEY field
+ * - Enforces limits on number and size of custom fields
+ */
 pub fn get_custom_fields_from_header(req: HttpRequest) -> HashMap<String, String> {

---

209-265: Good test coverage with various scenarios.

The tests cover key scenarios including:

• Custom fields extraction
• Ignored headers handling
• Format key special case
• Empty key validation

Consider adding two additional tests to verify limit handling:

1. A test that verifies truncation when field values exceed MAX_FIELD_VALUE_LENGTH
2. A test that verifies the maximum number of custom fields is enforced

#[test]
fn test_field_value_truncation() {
    let long_value = "a".repeat(MAX_FIELD_VALUE_LENGTH + 50);
    let req = TestRequest::default()
        .insert_header(("x-p-long", long_value.as_str()))
        .to_http_request();

    let custom_fields = get_custom_fields_from_header(req);
    
    assert_eq!(custom_fields.get("long").unwrap().len(), MAX_FIELD_VALUE_LENGTH);
}

#[test]
fn test_max_fields_limit() {
    let mut req_builder = TestRequest::default();
    
    // Add more than MAX_CUSTOM_FIELDS custom headers
    for i in 0..MAX_CUSTOM_FIELDS + 5 {
        req_builder = req_builder.insert_header((format!("x-p-field{}", i), "test"));
    }
    
    let req = req_builder.to_http_request();
    let custom_fields = get_custom_fields_from_header(req);
    
    // We expect MAX_CUSTOM_FIELDS (discounting the default fields that are always added)
    assert!(custom_fields.len() <= MAX_CUSTOM_FIELDS);
}

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro